package com.hk.core.elasticsearch.repository;

import co.elastic.clients.elasticsearch.core.search.FieldSuggester;
import co.elastic.clients.elasticsearch.core.search.Suggester;
import com.hk.commons.query.Order;
import com.hk.commons.query.QueryModel;
import com.hk.commons.query.SortUtils;
import com.hk.commons.query.page.QueryPage;
import com.hk.commons.query.page.SimpleQueryPage;
import com.hk.commons.util.BeanUtils;
import com.hk.commons.util.CollectionUtils;
import com.hk.core.elasticsearch.query.Condition;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Persistable;
import org.springframework.data.elasticsearch.client.elc.NativeQuery;
import org.springframework.data.elasticsearch.client.elc.NativeQueryBuilder;
import org.springframework.data.elasticsearch.core.ReactiveElasticsearchOperations;
import org.springframework.data.elasticsearch.core.SearchHit;
import org.springframework.data.elasticsearch.core.document.Document;
import org.springframework.data.elasticsearch.core.query.Criteria;
import org.springframework.data.elasticsearch.core.query.CriteriaQuery;
import org.springframework.data.elasticsearch.core.query.UpdateQuery;
import org.springframework.data.elasticsearch.core.suggest.response.Suggest;
import org.springframework.data.elasticsearch.repository.support.ElasticsearchEntityInformation;
import org.springframework.data.elasticsearch.repository.support.SimpleReactiveElasticsearchRepository;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Objects;

/**
 * 自定义扩展 Reactive Repository, 添加自定义查询方法
 *
 * @author Kevin
 * @date 2022/11/26 9:00
 */
public class SimpleBaseReactiveElasticsearchRepository<T extends Persistable<ID>, ID extends Serializable>
        extends SimpleReactiveElasticsearchRepository<T, ID>
        implements BaseReactiveElasticsearchRepository<T, ID> {

    private final ReactiveElasticsearchOperations reactiveElasticsearchOperations;

    private final ElasticsearchEntityInformation<T, ID> entityInformation;

    public SimpleBaseReactiveElasticsearchRepository(ElasticsearchEntityInformation<T, ID> entityInformation,
                                                     ReactiveElasticsearchOperations operations) {
        super(entityInformation, operations);
        this.entityInformation = entityInformation;
        this.reactiveElasticsearchOperations = operations;
    }

    @Override
    public Mono<QueryPage<T>> queryForPage(QueryModel<T> queryModel) {
        var map = BeanUtils.beanToMap(queryModel.getParam());
        var criteria = new Criteria();
        if (CollectionUtils.isNotEmpty(map)) {
            map.forEach((key, value) -> criteria.and(Criteria.where(key).is(value)));
        }
        var query = new CriteriaQuery(criteria, PageRequest.of(queryModel.getStartRowIndex(),
                queryModel.getPageSize(), SortUtils.toSort(queryModel.getOrders())));
        return reactiveElasticsearchOperations.searchForHits(query,
                        entityInformation.getJavaType())
                .flatMap(hits -> {
                    long totalHits = hits.getTotalHits();
                    return hits.getSearchHits()
                            .map(SearchHit::getContent)
                            .collectList()
                            .map(contents -> new SimpleQueryPage<>(queryModel, contents, totalHits));
                });
    }

    @Override
    public Mono<QueryPage<T>> queryForPage(List<Condition> conditions, int pageIndex,
                                           int pageSize, Order... orders) {
        var query = new CriteriaQuery(new Criteria(),
                PageRequest.of(pageIndex, pageSize, SortUtils.toSort(orders)));
        Condition.addCriteria(query, conditions);
        return reactiveElasticsearchOperations.searchForHits(query, entityInformation.getJavaType())
                .flatMap(hits -> {
                    long totalHits = hits.getTotalHits();
                    return hits.getSearchHits()
                            .map(SearchHit::getContent)
                            .collectList()
                            .map(contents -> new SimpleQueryPage<>(contents, totalHits, pageIndex, pageSize));
                });
    }

    @Override
    public Mono<Long> count(T t) {
        var criteria = new Criteria();
        var map = BeanUtils.beanToMap(t);
        if (CollectionUtils.isNotEmpty(map)) {
            map.forEach((key, value) -> criteria.and(Criteria.where(key).is(value)));
        }
        return reactiveElasticsearchOperations.count(new CriteriaQuery(criteria),
                entityInformation.getJavaType());
    }

    @Override
    public Flux<T> findAll(List<Condition> conditions, Order... orders) {
        var query = new CriteriaQuery(new Criteria());
        Condition.addCriteria(query, conditions);
        query.addSort(SortUtils.toSort(orders));
        return reactiveElasticsearchOperations.search(query, entityInformation.getJavaType())
                .map(SearchHit::getContent);
    }

    @Override
    public Mono<Void> bulkUpdate(Collection<T> entities) {
        if (CollectionUtils.isNotEmpty(entities)) {
            var list = new ArrayList<UpdateQuery>(entities.size());
            entities.forEach(item -> {
                var id = item.getId();
                if (Objects.nonNull(id)) {
                    var updateMap = BeanUtils.beanToMap(item, "class", "new", "id");
                    if (CollectionUtils.isNotEmpty(updateMap)) {
                        list.add((UpdateQuery.builder(String.valueOf(id))).withDocument(Document.from(updateMap)).build());
                    }
                }
            });
            return reactiveElasticsearchOperations.bulkUpdate(list, entityInformation.getIndexCoordinates());
        }
        return Mono.empty();
    }

    @Override
    public Flux<String> suggest(FieldSuggester suggester) {
        var suggestionName = "suggest";
        var query = new NativeQuery(new NativeQueryBuilder()
                .withSuggester(Suggester.of(builder -> builder.suggesters(suggestionName, suggester))));
        return reactiveElasticsearchOperations.searchForHits(query, entityInformation.getJavaType())
                .flatMapMany(hits -> {
                    var suggest = hits.getSuggest();
                    if (Objects.isNull(suggest)) {
                        return Flux.empty();
                    }
                    var suggestion = suggest.getSuggestion(suggestionName);
                    return Flux.fromStream(suggestion.getEntries().stream()
                            .map(Suggest.Suggestion.Entry::getText));
                });
    }
}
